import requests
import os
import time
import re
from downloader_ui import TaskMan

_URL_RE = re.compile(r'http.*')
session = requests.Session()


def _get(*args, **kwargs):
    for i in range(10):
        try:
            resp = session.get(*args, **kwargs)
        except Exception as _:
            time.sleep(2)
            continue
        if resp.status_code not in (200, 206):
            time.sleep(1)
            continue
        return resp
    return None


def _download_m3u8(url: str, fp, episode: str, task_man: TaskMan) -> bool:
    # 1. 获取 m3u8 列表文件
    resp = _get(url)
    if not resp:
        return False

    urls = _URL_RE.findall(resp.text)
    if not urls:
        return False

    task_man.new_task(episode, total=len(urls))
    # 开始下载
    try:
        for url in urls:
            url = url.strip()
            resp = _get(url)
            if not resp:
                raise Exception()
            data = resp.content
            if data[:4] == b'\x89\x50\x4e\x47':
                fp.write(data[120:])
            else:
                fp.write(data)
            task_man.update_task(episode, 1)
    except Exception as e:
        return False

    return True


def _download_mp4(url: str, fp, episode: str, task_man: TaskMan) -> bool:
    i = 0
    chunk_size = 1024 * 1024

    def get_more():
        return _get(url, headers={'range': f'bytes={i}-{i + chunk_size - 1}'})

    try:
        resp = get_more()
        total = resp.headers.get('Content-Range', "0")
        it = total.rfind('/') + 1
        total = int(total[it:])
        task_man.new_task(episode, total=total, useBytes=True)

        while True:
            fp.write(resp.content)
            data_size = len(resp.content)
            i += data_size
            task_man.update_task(episode, data_size)
            if i >= total:
                break
            resp = get_more()
    except Exception as _:
        return False

    return True


def download(url: str, name: str, episode: str, task_man: TaskMan):
    # 创建保存剧集的文件夹
    try:
        if not os.path.isdir(name):
            os.makedirs(name)
    except Exception as e:
        return None
    # 定义文件名
    ext = os.path.splitext(url)[1]
    if ext == '.m3u8':
        ext = '.ts'
    else:
        ext = '.mp4'
    fname = os.path.abspath(os.path.join(name, episode + ext))

    ok = False
    with open(fname, 'wb') as fp:
        if ext == '.ts':
            ok = _download_m3u8(url, fp, episode, task_man)
        elif ext == '.mp4':
            ok = _download_mp4(url, fp, episode, task_man)
        else:
            return None
    return fname if ok else None
